package net.tp.algo.graph.mst;

import net.tp.algo.graph.Graph;
import net.tp.algo.graph.UndirectedGraph;
import net.tp.algo.heap.BinaryHeap;

import java.util.Arrays;
import java.util.HashSet;
import java.util.Set;

/**
 *
 * @author Trung Phan
 */
public class Prim {

    /**
     *
     * @param graph undirected weighted graph
     * @return minimum spanning tree (expressed as graph)
     */
    public static Graph mst(Graph graph) {

        int N = graph.order();

        UndirectedGraph result = new UndirectedGraph(N);

        Set<Integer> settled = new HashSet<>();

        Graph.Edge[] edgeTo = new Graph.Edge[N];

        // distance from the settled set to the node w, unlike Dijkstra's algorithm where distance is between the source and node w
        double[] distanceTo = new double[N];
        Arrays.fill(distanceTo, Double.MAX_VALUE);

        BinaryHeap<Integer> heap = new BinaryHeap<>( (u, w) -> Double.compare(distanceTo[u], distanceTo[w]) );

        for (int s : graph.vertices()) {

            if (settled.contains(s)) continue;
            heap.add(s);

            while (!heap.empty()) {
                int u = heap.removeHead();
                distanceTo[u] = 0;
                settled.add(u);

                if (edgeTo[u] != null) {
                    result.addEdge(u, edgeTo[u].other(u), edgeTo[u].weight());
                }

                for (Graph.Edge edge : graph.edges(u)) {
                    int w = edge.other(u);
                    if (!settled.contains(w)) {
                        if (distanceTo[w] > edge.weight()) {
                            distanceTo[w] = edge.weight();
                            edgeTo[w] = edge;
                            if (heap.contains(w)) heap.updatePriority(w);
                            else heap.add(w);
                        }
                    }
                }
            }
        }

        return result;
    }

}
